{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "# %pip install -qU langchain langchain-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "# os.environ[\"OPENAI_API_KEY\"] = getpass.getpass()\n",
    "\n",
    "# Optional, uncomment to trace runs with LangSmith. Sign up here: https://smith.langchain.com.\n",
    "# os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "# os.environ[\"LANGCHAIN_API_KEY\"] = getpass.getpass()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "from typing import Literal, Optional, Tuple\n",
    "\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "\n",
    "\n",
    "class SubQuery(BaseModel):\n",
    "    \"\"\"Search over a database of tutorial videos about a software library.\"\"\"\n",
    "\n",
    "    sub_query: str = Field(\n",
    "        ...,\n",
    "        description=\"A very specific query against the database.\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.output_parsers import PydanticToolsParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "system = \"\"\"You are an expert at converting user questions into database queries. \\\n",
    "You have access to a database of tutorial videos about a software library for building LLM-powered applications. \\\n",
    "\n",
    "Perform query decomposition. Given a user question, break it down into distinct sub questions that \\\n",
    "you need to answer in order to answer the original question.\n",
    "\n",
    "If there are acronyms or words you are not familiar with, do not try to rephrase them.\"\"\"\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
    "llm_with_tools = llm.bind_tools([SubQuery])\n",
    "parser = PydanticToolsParser(tools=[SubQuery])\n",
    "query_analyzer = prompt | llm_with_tools | parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[SubQuery(sub_query='How to do rag')]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_analyzer.invoke({\"question\": \"how to do rag\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[SubQuery(sub_query='How to use multi-modal models in a chain'),\n",
       " SubQuery(sub_query='How to turn a chain into a REST API')]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_analyzer.invoke(\n",
    "    {\n",
    "        \"question\": \"how to use multi-modal models in a chain and turn chain into a rest api\"\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "examples = []\n",
    "question = \"What's chat langchain, is it a langchain template?\"\n",
    "queries = [\n",
    "    SubQuery(sub_query=\"What is chat langchain\"),\n",
    "    SubQuery(sub_query=\"What is a langchain template\"),\n",
    "]\n",
    "examples.append({\"input\": question, \"tool_calls\": queries})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"How would I use LangGraph to build an automaton\"\n",
    "queries = [\n",
    "    SubQuery(sub_query=\"How to build automaton with LangGraph\"),\n",
    "]\n",
    "examples.append({\"input\": question, \"tool_calls\": queries})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"How to build multi-agent system and stream intermediate steps from it\"\n",
    "queries = [\n",
    "    SubQuery(sub_query=\"How to build multi-agent system\"),\n",
    "    SubQuery(sub_query=\"How to stream intermediate steps\"),\n",
    "    SubQuery(sub_query=\"How to stream intermediate steps from multi-agent system\"),\n",
    "]\n",
    "examples.append({\"input\": question, \"tool_calls\": queries})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"What's the difference between LangChain agents and LangGraph?\"\n",
    "queries = [\n",
    "    SubQuery(sub_query=\"What's the difference between LangChain agents and LangGraph?\"),\n",
    "    SubQuery(sub_query=\"What are LangChain agents\"),\n",
    "    SubQuery(sub_query=\"What is LangGraph\"),\n",
    "]\n",
    "examples.append({\"input\": question, \"tool_calls\": queries})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Is the teacher standing in front of the classroom while giving a lecture or conducting an activity?\"\n",
    "queries = [\n",
    "    SubQuery(sub_query=\"Is the teacher standing in front of the classroom?\"),\n",
    "    SubQuery(sub_query=\"Is the teacher giving a lecture or conducting an activity?\"),\n",
    "]\n",
    "\n",
    "examples.append({\"input\": question, \"tool_calls\": queries})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Update prompt templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "from typing import Dict, List\n",
    "\n",
    "from langchain_core.messages import (\n",
    "    AIMessage,\n",
    "    BaseMessage,\n",
    "    HumanMessage,\n",
    "    SystemMessage,\n",
    "    ToolMessage,\n",
    ")\n",
    "\n",
    "\n",
    "def tool_example_to_messages(example: Dict) -> List[BaseMessage]:\n",
    "    messages: List[BaseMessage] = [HumanMessage(content=example[\"input\"])]\n",
    "    openai_tool_calls = []\n",
    "    for tool_call in example[\"tool_calls\"]:\n",
    "        openai_tool_calls.append(\n",
    "            {\n",
    "                \"id\": str(uuid.uuid4()),\n",
    "                \"type\": \"function\",\n",
    "                \"function\": {\n",
    "                    \"name\": tool_call.__class__.__name__,\n",
    "                    \"arguments\": tool_call.json(),\n",
    "                },\n",
    "            }\n",
    "        )\n",
    "    messages.append(\n",
    "        AIMessage(content=\"\", additional_kwargs={\"tool_calls\": openai_tool_calls})\n",
    "    )\n",
    "    tool_outputs = example.get(\"tool_outputs\") or [\n",
    "        \"This is an example of a correct usage of this tool. Make sure to continue using the tool this way.\"\n",
    "    ] * len(openai_tool_calls)\n",
    "    for output, tool_call in zip(tool_outputs, openai_tool_calls):\n",
    "        messages.append(ToolMessage(content=output, tool_call_id=tool_call[\"id\"]))\n",
    "    return messages\n",
    "\n",
    "\n",
    "example_msgs = [msg for ex in examples for msg in tool_example_to_messages(ex)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import MessagesPlaceholder\n",
    "\n",
    "system = \"\"\"You are an expert at converting user questions into database queries. \\\n",
    "You have access to a database of tutorial videos about a software library for building LLM-powered applications. \\\n",
    "\n",
    "Perform query decomposition. Given a user question, break it down into the most specific sub questions you can \\\n",
    "which will help you answer the original question. Each sub question should be about a single concept/fact/idea.\n",
    "\n",
    "If there are acronyms or words you are not familiar with, do not try to rephrase them.\"\"\"\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        MessagesPlaceholder(\"examples\", optional=True),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "query_analyzer_with_examples = (\n",
    "    prompt.partial(examples=example_msgs) | llm_with_tools | parser\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[SubQuery(sub_query=\"What's the difference between Web Voyager and Reflection Agents?\"),\n",
       " SubQuery(sub_query='Do Web Voyager and Reflection Agents use LangGraph?')]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_analyzer_with_examples.invoke(\n",
    "    {\n",
    "        \"question\": \"what's the difference between web voyager and reflection agents? do they use langgraph?\"\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
